// Fill out your copyright notice in the Description page of Project Settings.

#include "WorldEntityManager.h"
#include "drivers/WorldEntityManagerDriverFile.h"
#include "entities/WorldEntityReferencing.h"
#include <mutex>

struct FWorldEntityManager::Opaque
{
	FWorldEntityManager*			EntityManager;
	int								WorldId;
	int								Revision;
	std::int64_t					TmpId;
	std::mutex						BigMutex;
	FWorldEntityManagerPool::Ptr	Pool;
	FWorldEntityManagerDriver::Ptr	Driver;

	std::unordered_set<FWorldEntity::Ptr>	NewEntities;
	std::unordered_set<FWorldEntity::Ptr>	DirtyEntities;
	std::unordered_set<FWorldEntity::Ptr>	RemovedEntities;

public:
	bool EnsureRevision()
	{
		// We should always operate on a specific version. We find the
		// latest revision in the get region request.
		if (Revision < 0)
		{
			Revision = Driver->GetRevision(WorldId);
		}
		return Revision >= 0;
	}

	void Resolve(const FWorldEntity::Ptr& entity)
	{
		FWorldEntityReferencing referencing(EntityManager);
		entity->Resolve(referencing);
	}

	void Resolve()
	{
		// Resolve the references between entities.
		for (auto& entity : Pool->GetTypeMapping<FWorldEntity>())
		{
			Resolve(entity);
		}
	}

	void TurnInto(FWorldEntity::EntityState state)
	{
		// Turn all entities into specified state.
		for (auto& entity : Pool->GetTypeMapping<FWorldEntity>())
		{
			entity->ChangeTo(state, EntityManager);
		}
	}
};

FWorldEntityManager::FWorldEntityManager()
	: _op(new Opaque())
{
	Reset("http://localhost:8080");
}

FWorldEntityManager::FWorldEntityManager(const std::string& endpointOrFile)
	: _op(new Opaque())
{
	Reset(endpointOrFile);
}

FWorldEntityManager::~FWorldEntityManager()
{
	// Entities are no longer managed.
	_op->TurnInto(FWorldEntity::EntityState::Detached);
}

void FWorldEntityManager::Reset(const std::string& content)
{
	Clear();
	_op->EntityManager = this;
	if (content.compare(0, 4, "http") == 0) {
		//_op->Driver = std::make_shared<FWorldEntityManagerDriverDB>();
	}
	else
		_op->Driver = std::make_shared<FWorldEntityManagerDriverFile>();

	if (_op->Driver)
		_op->Driver->Connect(content);
}

int FWorldEntityManager::GetRevision() const
{
	// An entity manager should always operate on a specific revision globally.
	return _op->Revision;
}

const std::unordered_set<typename FWorldEntity::Ptr>& FWorldEntityManager::GetAll(const FWorldEntityMeta& meta) const
{
	return _op->Pool->TypeMappings[meta.ClassName];
}

typename FWorldEntity::Ptr FWorldEntityManager::GetById(const FWorldEntityMeta& meta, std::int64_t id) const
{
	// IDs in OpenDRIVE is not unique between classes. e.g. there can be both road #3 and junction #3.
	// ID alone is not able to distinguish the element type.
	auto& idmapping = _op->Pool->IdMappings[meta.ClassName];
	auto it = idmapping.find(id);
	if (it != idmapping.end())
	{
		if (it->second->MetaClass().IsA(meta))
		{
			return it->second;
		}
	}

	return nullptr;
}

void FWorldEntityManager::Region(double E, double W, double N, double S, int world, int rev)
{
	// This method doesn't support merging. Only one region is supported.
	Clear();

	// Working on a different world ?
	_op->WorldId = world;

	// Find the appropriate revision to work on.
	if (!_op->EnsureRevision()) return;

	// Working on a specific revision ?
	if (rev >= 0)
	{
		_op->Revision = rev;
	}

	// Staging area
	FWorldEntityManagerPool::Ptr staging =
		_op->Driver->GetRegion(E, W, N, S, _op->WorldId, _op->Revision);

	// Commit if no exception.
	_op->Pool = std::move(staging);
	_op->Resolve();
	_op->TurnInto(FWorldEntity::EntityState::Managed);
}

void FWorldEntityManager::Clear()
{
	_op->Revision = -1;
	_op->TmpId = -100;
	_op->Pool = std::make_shared<FWorldEntityManagerPool>();
	_op->NewEntities.clear();
	_op->DirtyEntities.clear();
	_op->RemovedEntities.clear();
}

FWorldEntityHeader::Ptr FWorldEntityManager::Header() const
{
	return _op->Pool->Header;
}

void FWorldEntityManager::Persist(const FWorldEntity::Ptr& entity)
{
	std::unique_lock<std::mutex> lock(_op->BigMutex);
	if (entity)
	{
		if (entity->IsManagedBy(this))
		{
			// A managed entity. Just trigger the changed observer.
			lock.unlock(); // mutex is not recursive
			entity->EntityChanged();
		}
		else if (entity->Is(FWorldEntity::EntityState::New))
		{
			// A new entity. Manage it. At this time, we don't know the ID
			// of the entity so assign a transient negative ID.
			entity->ID.Get() = _op->TmpId--;
			_op->Pool->Add(entity);
			_op->NewEntities.insert(entity);
			entity->ChangeTo(FWorldEntity::EntityState::Managed, this);
		}
	}
}

void FWorldEntityManager::Remove(const FWorldEntity::Ptr& entity)
{
	std::unique_lock<std::mutex> lock(_op->BigMutex);
	if (entity && entity->IsManagedBy(this))
	{
		// Move to "deleted" area so that the entity is not visible when
		// the entity is not pending add.
		if (_op->NewEntities.find(entity) == _op->NewEntities.end())
		{
			_op->RemovedEntities.insert(entity);
		}

		_op->Pool->Remove(entity);
		_op->NewEntities.erase(entity);
		_op->DirtyEntities.erase(entity);
		entity->ChangeTo(FWorldEntity::EntityState::Removed, this);
	}
}

void FWorldEntityManager::Flush()
{
	// Grab the changes until now.
	std::unordered_set<FWorldEntity::Ptr> newEntities;
	std::unordered_set<FWorldEntity::Ptr> dirtyEntities;
	std::unordered_set<FWorldEntity::Ptr> removedEntities;
	{
		std::unique_lock<std::mutex> lock(_op->BigMutex);
		std::swap(newEntities, _op->NewEntities);
		std::swap(dirtyEntities, _op->DirtyEntities);
		std::swap(removedEntities, _op->RemovedEntities);
		for (const auto& entity : newEntities) _op->Pool->Remove(entity);
	}

	// Make changes.
	try
	{
		_op->Driver->SetRegion(_op->Pool, newEntities, dirtyEntities, removedEntities, _op->WorldId, _op->Revision);
	}
	catch (std::exception&)
	{
		// Rollback changes.
		std::unique_lock<std::mutex> lock(_op->BigMutex);
		_op->NewEntities.insert(newEntities.begin(), newEntities.end());
		_op->DirtyEntities.insert(dirtyEntities.begin(), dirtyEntities.end());
		_op->RemovedEntities.insert(removedEntities.begin(), removedEntities.end());
		for (const auto& entity : newEntities) _op->Pool->Add(entity);
		throw;
	}

	// Survive!
	_op->Revision++; // no exception is thrown
	_op->TmpId = -100;

	// Add
	for (auto& entity : newEntities)
	{
		entity->EntityFlushed();
	}

	// Edit
	for (auto& entity : dirtyEntities)
	{
		entity->EntityFlushed();
	}

	// Remove
	for (auto& entity : removedEntities)
	{
		entity->EntityFlushed();
		entity->ChangeTo(FWorldEntity::EntityState::Detached, this); // no longer managed
	}

	// Update the pool for the new IDs.
	for (auto& entity : newEntities)
	{
		_op->Pool->Add(entity);
	}

	// Resolve references for transient IDs.
	for (auto& entity : newEntities)
	{
		_op->Resolve(entity);
	}
	for (auto& entity : dirtyEntities)
	{
		_op->Resolve(entity);
	}
}

void FWorldEntityManager::EntityChanged(const FWorldEntity::Ptr& entity)
{
	std::unique_lock<std::mutex> lock(_op->BigMutex);
	if (entity)
	{
		// An entity is dirty when it's not pending add or remove.
		if (_op->NewEntities.find(entity) == _op->NewEntities.end() &&
			_op->RemovedEntities.find(entity) == _op->RemovedEntities.end())
		{
			_op->DirtyEntities.insert(entity);
		}
	}
}
